# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import json
import logging
from typing import List, Union

from tornado import gen

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.coinone_v2.kr_rest.product import CoinoneProduct
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase
from coin.exchange.coinone_v2.ws.config import orderbook_unit_map, group_symbol


class CoinoneSubscriber(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._symbol = []
    for product in to_list(symbol):
      self._symbol.append(CoinoneProduct.FromStr(product).native_subscription_symbol)
    assert len(self._symbol) == 1, 'Coinone orderbook channel only supports single symbol!'

    self._register_url('wss://wss.coinone.co.kr/ws?token_type=web')
    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_depth = topic_map.search_topic_id('coinone_ws_depth_%s%s' % (group_id, worker_id))[0]
    topic_trade = topic_map.search_topic_id('coinone_ws_trade_%s%s' % (group_id, worker_id))[0]

    self._publisher = {
        'orderbook': queue.get_writer(topic_id=topic_depth[1]),
        'trades': queue.get_writer(topic_id=topic_trade[1]),
    }

    self._enable_topic_event_publisher(queue, 'coinone_ws_topic_event', [topic_depth[0]])
    self._add_teardown_callback(self._teardown_publisher)

  @gen.coroutine
  def on_open(self):
    yield self._send_subscribe_request()

  @gen.coroutine
  def _send_subscribe_request(self):
    for subscribe_symbol in self._symbol:
      subscribe_depth_str = json.dumps({
          "event": "subscribe",
          "channel": "orderbook",
          "trading_pair": subscribe_symbol,
          "unit": orderbook_unit_map[subscribe_symbol]
      })
      yield self._ws.write_message(subscribe_depth_str)

      subscribe_trade_str = json.dumps({
          "event": "subscribe", "channel": "trades", "trading_pairs": [subscribe_symbol]
      })
      yield self._ws.write_message(subscribe_trade_str)

      yield gen.sleep(0.5)  # Sleep 0.5 second before next request

    self._publish_topic_event('START')
    self._logger.info('Subscribe done')

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def on_message(self, timestamp: int, message: bytes):
    if self.stopped:
      return

    event = message.get('event', None)
    if event is not None:
      if event == 'subscribed':
        self._publish_topic_event(message)
      else:
        self._logger.error('Unknown message: %s' % message)
    else:
      try:
        channel = message['channel']
        self._publisher[channel].write(timestamp, message)
      except (KeyError, ValueError):
        self._logger.error('Unknown channel: %s' % message)


class CoinoneMultiSubscriberWrapper(object):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbols: Union[List[str], str],
               close_callback=None,
               logger=None):
    assert group_id == "TBD"
    self._subscribers = []
    self._symbols = []
    for product in to_list(symbols):
      self._symbols.append(CoinoneProduct.FromStr(product).native_subscription_symbol)

    for symbol in self._symbols:
      group_id = None
      for k, v in group_symbol.items():
        if symbol in v:
          group_id = k

      if group_id is None:
        logging.info("Unknown symbol: %s, %s. Dynamically resolved!", symbol, group_id)
        group_id = 'a'

      self._subscribers.append(
          CoinoneSubscriber(queue, group_id, worker_id, symbol, close_callback, logger))

  def subscribe(self):
    for subscriber in self._subscribers:
      subscriber.subscribe()

  def start(self, *args, **kwargs):
    for subscriber in self._subscribers:
      subscriber.start(*args, **kwargs)

  def stop(self, *args, **kwargs):
    for subscriber in self._subscribers:
      subscriber.stop(*args, **kwargs)
